/**
 * The MIT License (MIT)
 *
 * Copyright (c) 2021 Marcus Britanicus (https://gitlab.com/marcusbritanicus)
 * Copyright (c) 2021 Abrar (https://gitlab.com/s96Abrar)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 **/

#include <QDebug>
#include <QtWidgets>

#include <wayland-client.h>
#include <wayqt/WayQtUtils.hpp>
#include <wayqt/Registry.hpp>
#include <wayqt/ScreenCopy.hpp>


QMap<wl_shm_format, QImage::Format> mFormats{
    { WL_SHM_FORMAT_XRGB8888, QImage::Format_ARGB32 },
    { WL_SHM_FORMAT_ARGB8888, QImage::Format_ARGB32 },
    { WL_SHM_FORMAT_XBGR8888, QImage::Format_RGBA8888 },
    { WL_SHM_FORMAT_ABGR8888, QImage::Format_RGBA8888 },
};

int main( int argc, char *argv[] ) {
    QApplication *app = new QApplication( argc, argv );

    WQt::Registry *reg = new WQt::Registry( WQt::Wayland::display() );

    if ( reg == nullptr ) {
        qCritical() << "Unable to obtain the Wayland Registry Object.";
        return 1;
    }

    reg->setup();

    /** Wait upto 5s for the interface to be up. */
    qInfo() << "Waiting for interface";

    if ( reg->waitForInterface( WQt::Registry::ScreenCopyManagerInterface, 5000 ) ) {
        WQt::ScreenCopyManager *scm = reg->screenCopier();

        if ( scm == nullptr ) {
            return 1;
        }

        WQt::ScreenCopyFrame *frame = scm->captureOutput( false, app->primaryScreen() );

        if ( frame == nullptr ) {
            return 1;
        }

        WQt::ScreenFrameBuffer *buffer = new WQt::ScreenFrameBuffer;

        QObject::connect(
            frame, &WQt::ScreenCopyFrame::bufferDone, [ = ] () {
                for ( WQt::FrameBufferInfo info: frame->availableFormats() ) {
                    if ( mFormats.contains( info.format ) ) {
                        buffer->info = info;
                        buffer->initializeBuffer( info );

                        if ( buffer->buffer == NULL ) {
                            qWarning() << "Failed to create buffer with format" << buffer->info.format;
                            continue;
                        }

                        frame->attachBuffer( buffer );

                        frame->copy();
                    }
                }
            }
        );

        QObject::connect(
            frame, &WQt::ScreenCopyFrame::ready, [ = ] ( WQt::ScreenFrameBuffer *buffer ) {
                if ( buffer == nullptr ) {
                    qDebug() << "Screenshot failed";
                }

                QImage img(
                    (uchar *)buffer->data,
                    buffer->info.width,
                    buffer->info.height,
                    buffer->info.stride,
                    mFormats[ buffer->info.format ]
                );

                if ( img.save( ( (argc == 1)? "image.png" : argv[ 1 ]) ) ) {
                    qDebug() << "Screenshot successful";
                }

                else {
                    qDebug() << "Screenshot failed";
                }

                app->quit();
            }
        );

        frame->setup();
    }

    return app->exec();
}
